package com.gwb.util;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.http.ContentType;
import cn.hutool.http.Header;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.Method;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import lombok.NonNull;
import okhttp3.*;


import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.File;
import java.io.PrintWriter;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 聊天接口流式 返回
 */
public class kimiUtils {

    private static final String API_KEY = "";
    private static final String MODELS_URL = "https://api.moonshot.cn/v1/models";
    private static final String FILES_URL = "https://api.moonshot.cn/v1/files";
    private static final String ESTIMATE_TOKEN_COUNT_URL = "https://api.moonshot.cn/v1/tokenizers/estimate-token-count";
    private static final String CHAT_COMPLETION_URL = "https://api.moonshot.cn/v1/chat/completions";

    public static String getModelList() {
        return getCommonRequest(MODELS_URL)
                .execute()
                .body();
    }

    /***
     * 上传文件
     * {"id":"cq99ijqtnn0t879b57fg","object":"file","bytes":133,
     *   "created_at":1720883535,"filename":"用户增长表.csv",
     *   purpose":"file-extract","status":"ok","status_details":""}
     * @param file
     * @return
     */
    public static String uploadFile(@NonNull File file) {
        return getCommonRequest(FILES_URL)
                .method(Method.POST)
                .header("purpose", "file-extract")
                .form("file", file)
                .execute()
                .body();
    }

    /***
     * 获取文件列表
     * {"object":"list","data":[{"id":"cq99ijqtnn0t879b57fg",
     * "object":"file","bytes":133,"created_at":1720883535,
     * "filename":"用户增长表.csv","purpose":"file-extract",
     * "status":"ok","status_details":""}]}
     * @return
     */
    public static String getFileList() {
        return getCommonRequest(FILES_URL)
                .execute()
                .body();
    }

    /***
     * 删除文件
     * {"deleted":true,"id":"cq99ijqtnn0t879b57fg","object":"file"}
     * @param fileId
     * @return
     */
    public static String deleteFile(@NonNull String fileId) {
        return getCommonRequest(FILES_URL + "/" + fileId)
                .method(Method.DELETE)
                .execute()
                .body();
    }

    /***
     * 获取文件信息详情
     * {"id":"cq99ijqtnn0t879b57fg","object":"file","bytes":133,
     * "created_at":1720883535,"filename":"用户增长表.csv",
     * "purpose":"file-extract","status":"ok","status_details":""}
     * @param fileId
     * @return
     */
    public static String getFileDetail(@NonNull String fileId) {
        return getCommonRequest(FILES_URL + "/" + fileId)
                .execute()
                .body();
    }

    /**
     * 获取文件内容
     * {"content":"时间: 2024.01.01\n增长用户数: 1\n: 时间: 2024.01.02\n
     * 增长用户数: 2\n: 时间: 2024.01.03\n增长用户数: 33\n: 时间: 2024.01.04\n
     * 增长用户数: 2\n: 时间: 2024.01.05\n增长用户数: 1\n: 时间: 2024.01.06\n
     * 增长用户数: 12\n: 时间: 2024.01.07\n增长用户数: 65\n: ",
     * "file_type":"application/vnd.ms-excel","filename":"用户增长表.csv","title":"","type":"file"}
     * @param fileId
     * @return
     */
    public static String getFileContent(@NonNull String fileId) {
        return getCommonRequest(FILES_URL + "/" + fileId + "/content")
                .execute()
                .body();
    }

    public static String estimateTokenCount(@NonNull String model, @NonNull List<Message> messages) {
        String requestBody = new JSONObject()
                .putOpt("model", model)
                .putOpt("messages", messages)
                .toString();
        return getCommonRequest(ESTIMATE_TOKEN_COUNT_URL)
                .method(Method.POST)
                .header(Header.CONTENT_TYPE, ContentType.JSON.getValue())
                .body(requestBody)
                .execute()
                .body();
    }

    public static String chat(@NonNull String model, @NonNull List<Message> messages, HttpServletResponse response) {
        // 需要指定response的ContentType为流式输出，且字符编码为UTF-8
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        // 禁用缓存
        response.setHeader("Cache-Control", "no-cache");
        BufferedReader reader =null;
        try(PrintWriter out = response.getWriter()){
            KimiClient kimiClient = (KimiClient) SpringUtil.getBean("kimiClient");
            String apiKey = kimiClient.getApiKey();
            String requestBody = new JSONObject()
                    .putOpt("model", model)
                    .putOpt("messages", messages)
                    .putOpt("stream", true)
                    .toString();
            Request okhttpRequest = new Request.Builder()
                    .url(CHAT_COMPLETION_URL)
                    .post(RequestBody.create(MediaType.get(ContentType.JSON.getValue()),requestBody))//防止第三方引入的时候报错
                    /*.post(RequestBody.create(requestBody, MediaType.get(ContentType.JSON.getValue())))*/
                    .addHeader("Authorization", "Bearer " + apiKey)
                    .build();
            Call call = new OkHttpClient().newCall(okhttpRequest);
            Response okhttpResponse = call.execute();
             reader = new BufferedReader(okhttpResponse.body().charStream());
            String line;
            AtomicReference<String> result = new AtomicReference<>("");
            while ((line = reader.readLine()) != null) {
                if (StrUtil.isBlank(line)) {
                    continue;
                }
                if (JSONUtil.isTypeJSON(line)) {
                    Optional.of(JSONUtil.parseObj(line))
                            .map(x -> x.getJSONObject("error"))
                            .map(x -> x.getStr("message"))
                            .ifPresent(x ->
                                    {
                                        System.out.println("error: " + x);
                                        //出问题抛异常
                                        throw new RuntimeException(x);
                                    }
                            );
                }
                line = StrUtil.replace(line, "data: ", StrUtil.EMPTY);
                if (StrUtil.equals("[DONE]", line) || !JSONUtil.isTypeJSON(line)) {
                    //解释返回结果
                    return result.get();
                }
                Optional.of(JSONUtil.parseObj(line))
                        .map(x -> x.getJSONArray("choices"))
                        .filter(CollUtil::isNotEmpty)
                        .map(x -> (JSONObject) x.get(0))
                        .map(x -> x.getJSONObject("delta"))
                        .map(x -> x.getStr("content"))
                        .ifPresent(x -> {
                            //后面往响应体中写入
                            result.set(result.get() + x);
                            try {
                                // 发送数据到前端
                                out.print(x);
                                out.flush(); // 确保数据立即发送
                            } catch (Exception e) {
                                e.printStackTrace();
                                throw new RuntimeException(e);
                            }
                            System.out.println("======== " + x);
                        });
            }
        }catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException(e);
        }finally {
            if (reader != null){
                try {
                    reader.close();
                }catch (Exception e){
                    e.printStackTrace();
                    throw new RuntimeException(e);
                }
            }
        }
        throw new RuntimeException("ai无回复");
    }

    public static String doChat(@NonNull String model, @NonNull List<Message> messages) {

        // 禁用缓存
        BufferedReader reader =null;
        try{
            KimiClient kimiClient = (KimiClient) SpringUtil.getBean("kimiClient");
            String apiKey = kimiClient.getApiKey();
            String requestBody = new JSONObject()
                    .putOpt("model", model)
                    .putOpt("messages", messages)
                    .putOpt("stream", true)
                    .toString();
            Request okhttpRequest = new Request.Builder()
                    .url(CHAT_COMPLETION_URL)
                    .post(RequestBody.create(MediaType.get(ContentType.JSON.getValue()),requestBody))//防止第三方引入的时候报错
                    /*.post(RequestBody.create(requestBody, MediaType.get(ContentType.JSON.getValue())))*/
                    .addHeader("Authorization", "Bearer " + apiKey)
                    .build();
            Call call = new OkHttpClient().newCall(okhttpRequest);
            Response okhttpResponse = call.execute();
            reader = new BufferedReader(okhttpResponse.body().charStream());
            String line;
            AtomicReference<String> result = new AtomicReference<>("");
            while ((line = reader.readLine()) != null) {
                if (StrUtil.isBlank(line)) {
                    continue;
                }
                if (JSONUtil.isTypeJSON(line)) {
                    Optional.of(JSONUtil.parseObj(line))
                            .map(x -> x.getJSONObject("error"))
                            .map(x -> x.getStr("message"))
                            .ifPresent(x ->
                                    {
                                        System.out.println("error: " + x);
                                        //出问题抛异常
                                        throw new RuntimeException(x);
                                    }
                            );
                }
                line = StrUtil.replace(line, "data: ", StrUtil.EMPTY);
                if (StrUtil.equals("[DONE]", line) || !JSONUtil.isTypeJSON(line)) {
                    //解释返回结果
                    return result.get();
                }
                Optional.of(JSONUtil.parseObj(line))
                        .map(x -> x.getJSONArray("choices"))
                        .filter(CollUtil::isNotEmpty)
                        .map(x -> (JSONObject) x.get(0))
                        .map(x -> x.getJSONObject("delta"))
                        .map(x -> x.getStr("content"))
                        .ifPresent(x -> {
                            //后面往响应体中写入
                            result.set(result.get() + x);
                            System.out.println("======== " + x);
                        });
            }
        }catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException(e);
        }finally {
            if (reader != null){
                try {
                    reader.close();
                }catch (Exception e){
                    e.printStackTrace();
                    throw new RuntimeException(e);
                }
            }
        }
        throw new RuntimeException("ai无回复");
    }


    private static HttpRequest getCommonRequest(@NonNull String url) {
        KimiClient kimiClient = (KimiClient) SpringUtil.getBean("kimiClient");
        String apiKey = kimiClient.getApiKey();
        return HttpRequest.of(url).header(Header.AUTHORIZATION, "Bearer " + apiKey);
    }

}
